import sys
import os
from PyQt6.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout,
                             QPushButton, QLabel, QFileDialog, QWidget, QMessageBox,
                             QTextEdit, QSplitter, QFrame)
from PyQt6.QtCore import Qt
from PyQt6.QtGui import QPixmap, QImage
import cv2
import numpy as np
from core.video_thread import VideoThread
from core.yolo_processor import YOLOProcessor
from utils.image_utils import ImageUtils
from resources.styles import Styles


class YOLOv8Deploy(QMainWindow):
    def __init__(self):
        super().__init__()
        self.model_processor = None
        self.video_thread = None
        self.current_image = None
        self.current_results = None
        self.initUI()

    def initUI(self):
        self.setWindowTitle('YOLOv8 模型部署 - 目标检测系统')
        self.setGeometry(100, 100, 1400, 900)

        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 主布局
        main_layout = QVBoxLayout()
        central_widget.setLayout(main_layout)

        # 标题
        title_label = QLabel('YOLOv8 目标检测系统')
        title_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        title_label.setStyleSheet(Styles.TITLE_STYLE)
        main_layout.addWidget(title_label)

        # 按钮布局
        button_layout = QHBoxLayout()

        # 加载模型按钮
        self.load_model_btn = QPushButton('加载YOLOv8模型')
        self.load_model_btn.clicked.connect(self.load_model)
        self.load_model_btn.setStyleSheet(Styles.BUTTON_LOAD_STYLE)
        button_layout.addWidget(self.load_model_btn)

        # 图片推理按钮
        self.image_btn = QPushButton('图片推理')
        self.image_btn.clicked.connect(self.process_image)
        self.image_btn.setStyleSheet(Styles.BUTTON_IMAGE_STYLE)
        self.image_btn.setEnabled(False)
        button_layout.addWidget(self.image_btn)

        # 视频推理按钮
        self.video_btn = QPushButton('视频推理')
        self.video_btn.clicked.connect(self.process_video)
        self.video_btn.setStyleSheet(Styles.BUTTON_VIDEO_STYLE)
        self.video_btn.setEnabled(False)
        button_layout.addWidget(self.video_btn)

        # 停止视频按钮
        self.stop_video_btn = QPushButton('停止视频')
        self.stop_video_btn.clicked.connect(self.stop_video)
        self.stop_video_btn.setStyleSheet(Styles.BUTTON_STOP_STYLE)
        self.stop_video_btn.setEnabled(False)
        button_layout.addWidget(self.stop_video_btn)

        # 清空显示按钮
        self.clear_btn = QPushButton('清空显示')
        self.clear_btn.clicked.connect(self.clear_display)
        self.clear_btn.setStyleSheet(Styles.BUTTON_CLEAR_STYLE)
        button_layout.addWidget(self.clear_btn)

        main_layout.addLayout(button_layout)

        # 创建分割器来分隔图像显示和结果展示
        splitter = QSplitter(Qt.Orientation.Horizontal)
        main_layout.addWidget(splitter)

        # 左侧：图像显示区域
        image_frame = QFrame()
        image_layout = QVBoxLayout(image_frame)

        image_title = QLabel('检测结果可视化')
        image_title.setAlignment(Qt.AlignmentFlag.AlignCenter)
        image_title.setStyleSheet(Styles.SECTION_TITLE_STYLE)
        image_layout.addWidget(image_title)

        self.result_label = QLabel()
        self.result_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.result_label.setStyleSheet(Styles.RESULT_LABEL_STYLE)
        self.result_label.setText('推理结果将显示在这里')
        self.result_label.setMinimumSize(600, 500)
        image_layout.addWidget(self.result_label)

        # 右侧：分类结果展示区域
        results_frame = QFrame()
        results_layout = QVBoxLayout(results_frame)

        results_title = QLabel('检测结果统计')
        results_title.setAlignment(Qt.AlignmentFlag.AlignCenter)
        results_title.setStyleSheet(Styles.SECTION_TITLE_STYLE)
        results_layout.addWidget(results_title)

        # 统计信息显示
        self.stats_label = QLabel()
        self.stats_label.setStyleSheet(Styles.STATS_LABEL_STYLE)
        self.stats_label.setText('检测统计信息将显示在这里')
        self.stats_label.setWordWrap(True)
        results_layout.addWidget(self.stats_label)

        # 详细结果文本框
        detailed_title = QLabel('详细检测结果')
        detailed_title.setStyleSheet(Styles.DETAILED_TITLE_STYLE)
        results_layout.addWidget(detailed_title)

        self.results_text = QTextEdit()
        self.results_text.setStyleSheet(Styles.RESULTS_TEXT_STYLE)
        self.results_text.setReadOnly(True)
        self.results_text.setPlaceholderText('详细检测结果将显示在这里...')
        results_layout.addWidget(self.results_text)

        # 将左右两侧添加到分割器
        splitter.addWidget(image_frame)
        splitter.addWidget(results_frame)
        splitter.setSizes([800, 400])

        # 状态栏
        self.status_label = QLabel('请先加载YOLOv8模型 (.pt文件)')
        self.status_label.setStyleSheet(Styles.STATUS_LABEL_STYLE)
        main_layout.addWidget(self.status_label)

    def load_model(self):
        """加载YOLOv8模型"""
        try:
            file_path, _ = QFileDialog.getOpenFileName(
                self, '选择YOLOv8模型文件', '', 'PyTorch Model Files (*.pt)')

            if file_path:
                self.status_label.setText(f'正在加载模型: {os.path.basename(file_path)}...')
                QApplication.processEvents()  # 更新UI

                # 加载模型
                self.model_processor = YOLOProcessor(file_path)

                # 启用按钮
                self.image_btn.setEnabled(True)
                self.video_btn.setEnabled(True)

                self.status_label.setText(f'模型加载成功: {os.path.basename(file_path)}')

                # 显示模型信息
                class_names = list(self.model_processor.get_class_names().values())
                self.results_text.setText(f"""模型加载成功！
模型路径: {file_path}
类别数量: {len(class_names)}
检测类别: {', '.join(class_names)}

等待进行推理...""")

                QMessageBox.information(self, '成功', 'YOLOv8模型加载成功！')

        except Exception as e:
            QMessageBox.critical(self, '错误', f'加载模型失败: {str(e)}')
            self.status_label.setText('模型加载失败')

    def process_image(self):
        """处理图片推理"""
        if self.model_processor is None:
            QMessageBox.warning(self, '警告', '请先加载YOLOv8模型！')
            return

        try:
            file_path, _ = QFileDialog.getOpenFileName(
                self, '选择图片文件', '',
                'Image Files (*.jpg *.jpeg *.png *.bmp *.tiff)')

            if file_path:
                self.status_label.setText('正在进行图片推理...')
                QApplication.processEvents()

                # 使用YOLOv8进行推理
                results = self.model_processor.process_image(file_path)

                # 获取带标注的结果图片
                annotated_image = results[0].plot()

                # 转换BGR到RGB
                annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)

                # 保存当前图片和结果用于显示
                self.current_image = annotated_image_rgb
                self.current_results = results[0]

                # 显示结果图片
                self.display_image(annotated_image_rgb)

                # 显示分类结果
                self.display_detection_results(results[0])

                self.status_label.setText('图片推理完成')

        except Exception as e:
            QMessageBox.critical(self, '错误', f'图片推理失败: {str(e)}')
            self.status_label.setText('图片推理失败')

    def process_video(self):
        """处理视频推理"""
        if self.model_processor is None:
            QMessageBox.warning(self, '警告', '请先加载YOLOv8模型！')
            return

        try:
            file_path, _ = QFileDialog.getOpenFileName(
                self, '选择视频文件', '',
                'Video Files (*.mp4 *.avi *.mov *.mkv *.wmv)')

            if file_path:
                # 停止之前的视频线程
                if self.video_thread and self.video_thread.isRunning():
                    self.video_thread.stop()
                    self.video_thread.wait()

                # 创建新的视频线程
                self.video_thread = VideoThread(self.model_processor, file_path)
                self.video_thread.frame_signal.connect(self.display_video_frame)
                self.video_thread.results_signal.connect(self.update_video_results)
                self.video_thread.finished_signal.connect(self.video_finished)

                # 更新按钮状态
                self.stop_video_btn.setEnabled(True)
                self.video_btn.setEnabled(False)
                self.image_btn.setEnabled(False)

                self.status_label.setText('正在播放视频推理结果...')

                # 清空之前的结果
                self.results_text.clear()
                self.stats_label.setText('实时检测中...')

                # 开始视频线程
                self.video_thread.start()

        except Exception as e:
            QMessageBox.critical(self, '错误', f'视频推理失败: {str(e)}')
            self.status_label.setText('视频推理失败')

    def display_detection_results(self, results):
        """显示检测结果"""
        if results.boxes is None or len(results.boxes) == 0:
            self.stats_label.setText('未检测到目标')
            self.results_text.setText('未检测到任何目标物体')
            return

        # 使用model_processor处理结果
        stats_text, results_text = self.model_processor.format_detection_results(results)

        self.stats_label.setText(stats_text)
        self.results_text.setText(results_text)

    def update_video_results(self, detections):
        """更新视频推理的实时结果"""
        stats_text, results_text = self.model_processor.format_video_results(detections)
        self.stats_label.setText(stats_text)
        self.results_text.setText(results_text)

    def stop_video(self):
        """停止视频播放"""
        if self.video_thread and self.video_thread.isRunning():
            self.video_thread.stop()
            self.video_thread.wait()

        self.video_finished()

    def video_finished(self):
        """视频播放结束"""
        self.stop_video_btn.setEnabled(False)
        self.video_btn.setEnabled(True)
        self.image_btn.setEnabled(True)
        self.status_label.setText('视频播放结束')
        self.stats_label.setText('视频检测已结束')

    def display_image(self, image_array):
        """显示图片"""
        try:
            pixmap = ImageUtils.numpy_to_qpixmap(image_array)

            # 缩放图片以适应显示区域，保持宽高比
            label_size = self.result_label.size()
            scaled_pixmap = pixmap.scaled(
                label_size.width() - 20,
                label_size.height() - 20,
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )

            self.result_label.setPixmap(scaled_pixmap)

        except Exception as e:
            print(f"显示图片错误: {e}")

    def display_video_frame(self, frame):
        """显示视频帧"""
        try:
            # 转换BGR到RGB
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            self.display_image(frame_rgb)
        except Exception as e:
            print(f"显示视频帧错误: {e}")

    def clear_display(self):
        """清空显示"""
        self.result_label.clear()
        self.result_label.setText('推理结果将显示在这里')
        self.current_image = None
        self.current_results = None
        self.stats_label.setText('检测统计信息将显示在这里')
        self.results_text.clear()
        self.status_label.setText('显示已清空')
